Skip to content

feat: add DeepSeek V4 Flash GRPO support - Automodel path - #3595

Merged
sharonyu-115 merged 45 commits into
NVIDIA-NeMo:mainfrom
zpqiu:dsv4-automodel-r060-port
Sep 14, 2026
Merged

sharonyu-115 merged 45 commits into
NVIDIA-NeMo:mainfrom
zpqiu:dsv4-automodel-r060-port

Conversation

@zpqiu

@zpqiu zpqiu commented Aug 12, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds GRPO support for deepseek-ai/DeepSeek-V4-Flash(-0731) and deepseek-ai/DeepSeek-V4-Flash-Base using the AutoModel training backend and colocated vLLM generation.

The PR adds:

  • DeepSeek V4 tokenizer/chat rendering support.
  • Blockwise FP8 weight refit into vLLM, including the model-specific packed parameter layouts and layerwise reload handling.
  • fp8_ds_mla KV-cache configuration and validation.
  • A 16-node reference recipe, manual smoke test, user guide, unit tests, and reference training curves through step 50.

Why are these changes needed?

DeepSeek V4 tokenizer

The supported checkpoints do not provide a Jinja chat template that Hugging Face apply_chat_template() can use. DeepSeek V4 instead defines its chat formatting through a Python encoder. Because the NeMo RL driver relies on apply_chat_template(), this PR vendors vLLM 0.25.1's DeepSeek V4 encoder and adapts it to the Hugging Face tokenizer interface. With the corresponding NeMo RL and vLLM tokenizer modes enabled by the recipe, training and rollout use consistent rendering, including enable_thinking.

fp8_ds_mla

fp8_ds_mla is a vLLM runtime KV-cache format rather than a value selected by the checkpoints' config.json. DeepSeek V4 sparse MLA uses a packed block-scaled layout containing FP8 NoPE state, BF16 RoPE state, and embedded block scales. This differs from the traditional per-tensor FP8 KV caches handled by NeMo RL's generic k_scale/v_scale refit path. This PR accepts the format while excluding it from per-tensor KV-scale calculation and synchronization, and validates the DeepSeek V4 configuration to keep the driver and realized vLLM cache layout consistent.

FP8 refit

The validated recipe trains dequantized BF16 weights with AutoModel while the colocated vLLM rollout model uses blockwise FP8 weights. DeepSeek V4 has model-specific Hugging Face-to-vLLM name mappings, fused projections, shape-changing wo_a BMM post-processing, and packed RoutedExperts layouts. These prevent repeated refits from directly loading checkpoint-shaped tensors into the already realized vLLM runtime parameters. The DeepSeek V4-scoped mapping and vLLM native layerwise reload lifecycle restore checkpoint-loadable layouts before each refit and rebuild the kernel layouts afterward.

Issues

N/A

Usage

The validated reference configuration uses 16 nodes with 8 H100 GPUs per node:

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/grpo-deepseek-v4-flash-0731-16n8g-automodel-cp8ep128.yaml

See DeepSeek V4 Flash for supported models, configuration details, and limitations.

Validation

  • CI-tested revision: e874e79.
  • Nightly test wandb: https://wandb.ai/nvidia/nemo-rl-pr3595
  • Relevant tokenizer, FP8 quantization/refit, vLLM backend, AutoModel setup, and DTensor worker unit tests are included.
  • The 16-node TP1/CP8/EP128 AutoModel + TP8/EP1 vLLM recipe was validated end to end; the guide includes raw training and validation curves through step 50.

Before your PR is "Ready for review"

Pre checks:

  • Read and followed the Contributor guidelines.
  • Added the necessary tests.
  • Ran the relevant unit, lint, and end-to-end functional validation.
  • Added the model guide, recipe documentation, and reference curves.

Additional Information

The validated scope is DeepSeek V4 Flash 0731 and Flash Base with AutoModel training and vLLM generation. DeepSeek V4 Pro, Megatron training, SGLang generation, and training-side FP8 fake quantization are outside the scope of this PR.

@copy-pr-bot

copy-pr-bot Bot commented Aug 12, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Aug 12, 2026
@zpqiu zpqiu changed the title [feat] add DeepSeek V4 Flash support - Automodel path feat: add DeepSeek V4 Flash support - Automodel path Aug 12, 2026
@zpqiu
zpqiu force-pushed the dsv4-automodel-r060-port branch 2 times, most recently from 074a7fd to 0bcd08a Compare August 29, 2026 06:56
Comment thread nemo_rl/models/automodel/setup.py Resolved
Comment thread nemo_rl/algorithms/utils.py Resolved
Comment thread nemo_rl/data/deepseek_v4_tokenizer.py Resolved
Comment thread tests/unit/algorithms/test_utils.py Resolved
@zpqiu
zpqiu force-pushed the dsv4-automodel-r060-port branch 2 times, most recently from 5d0a67b to 4286bb0 Compare September 5, 2026 13:19
@zpqiu zpqiu added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Sep 5, 2026
@zpqiu

zpqiu commented Sep 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f9511e0

@zpqiu

zpqiu commented Sep 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 0e43437

@zpqiu
zpqiu force-pushed the dsv4-automodel-r060-port branch from d6a3cf1 to 2020329 Compare September 6, 2026 15:05
@zpqiu

zpqiu commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2020329

@zpqiu

zpqiu commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4ba8402

@zpqiu

zpqiu commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 879376f

@zpqiu
zpqiu force-pushed the dsv4-automodel-r060-port branch from 879376f to 0a4e04a Compare September 7, 2026 02:12
@zpqiu

zpqiu commented Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 0a4e04a

@zpqiu
zpqiu force-pushed the dsv4-automodel-r060-port branch from a0748c7 to a1d83ee Compare September 7, 2026 07:16

@zpqiu zpqiu left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed with a team of six agents (RL/codebase, vLLM, Automodel, tests, bug-hunting, and triage of the existing threads), against vLLM 0.25.1 (752a3a50) and Automodel 1814c6c9. Every finding below was re-verified against source before posting.

Nice work — this is a large, genuinely difficult integration and it holds up well under scrutiny. Highlights from what we checked and found clean:

  • nemo_rl/data/deepseek_v4_encoding.py is byte-identical to vLLM 0.25.1's copy apart from one added license line — no adaptation drift.
  • All three non-obvious vLLM APIs are used correctly: replace_parameter(..., prefer_copy=True), deepgemm_post_process_fp8_weight_block(..., is_bmm=, bmm_batch_size=), and the get_layerwise_info(...).restore_metadata / SKIP_TENSORS lifecycle. The ordering around record_metadata_for_reloading is the only window that works, and the PR hits it.
  • The TP1/CP8/EP128 recipe is valid and internally consistent (ep_shard_size = 1 is the correct branch at exact equality; EP128 divides n_routed_experts: 256).
  • Every DeepSeek V4 test genuinely executes — 0 skipped, confirmed three ways (item-count reconciliation, the runner's nodeid manifest, and a negative control showing the marker filter bites). ruff format, ruff check, pyrefly, and config_cli.py minimize-check are all clean.
  • The prior round's report of 21 failing committed tests does not reproduce at this HEAD; all 21 existing review threads are stale or fixed, and we've resolved them.

Two notes on scope, so the findings below read correctly. Several are latent rather than live: the rendering issue is not on this recipe's GRPO path, and the dp_size guard is pre-existing code your recipe is simply the first to make degenerate. They're raised because they'd bite the next person, not because they affect your validated run.

One request: would you be able to share throughput numbers for the run behind the reference curves — tokens/sec and peak memory at the 16-node TP1/CP8/EP128 configuration? The convergence curves through step 50 are great for correctness, but there's no baseline for what the FP8 refit path costs per step, which is what the next person sizing a DSV4 run will want.

Generated by Claude Code

Comment thread nemo_rl/data/deepseek_v4_tokenizer.py Resolved

@zpqiu zpqiu left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Follow-up to my earlier review with three findings that landed after I published it -- my mistake for posting before the whole team had reported back. Two of these are things I should have caught in the first pass.

  • data.filter_overlong_prompts: true is read by nothing. It appears exactly once in the repo (this recipe) and is not a DataConfig field. Same class as the dead NRL_* env vars already flagged.
  • deepseek_v4_fp8.py is missing from pyrefly.toml while both its quantization siblings are listed, which is the only reason pyrefly check reports clean. Adding it surfaces one real error. This qualifies my earlier "lint is clean at HEAD" note: it is clean, but the PR's largest new module is not being typechecked.
  • A stride-shaped hole in the DSV4 test suite, found by mutation testing rather than inspection. The production code is correct; the suite simply would not catch one specific silent regression. A verified test is included.

Everything else from the first review stands unchanged.

Generated by Claude Code

Comment thread pyrefly.toml Resolved

@zpqiu zpqiu left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Final follow-up -- four test-quality findings, closing out the review. Sorry for the trailing reviews; the test agent's mutation work finished after I'd already posted, and these were worth not dropping.

  • num_special_tokens_to_add returns 1 instead of 0 against a real HF tokenizer. This sharpens my earlier comment on the same area: the issue is not only that the test is weak, the production method is wrong and the stub masks it.
  • process_weights_after_loading_moe has zero test coverage repo-wide, and this PR changed how it manages parameter storage. The sibling linear path has an incident-driven guard for exactly this failure mode; the MoE path holds more memory and has none. Verified test included.
  • The narrowed grpo.setup() KV-cache guard is untested on either side of the change.
  • A finalize_refit negative test cannot fail -- confirmed by mutation.

All findings verified against source; both included tests were run at HEAD and mutation-tested to confirm they are not tautologies.

To restate the overall picture, since this is the third review and the volume may read worse than the PR deserves: the implementation holds up well. The vendored encoder is byte-identical to upstream, all three non-obvious vLLM APIs are used correctly, the CP8/EP128 parallelism is valid, and every DeepSeek V4 test genuinely executes with zero skips (proven four ways). Roughly as many candidate findings were killed under scrutiny as were posted -- including two of my own.

Generated by Claude Code

Comment thread nemo_rl/data/deepseek_v4_tokenizer.py Resolved
Comment thread nemo_rl/algorithms/grpo.py Resolved
@zpqiu zpqiu changed the title feat: add DeepSeek V4 Flash support - Automodel path feat: add DeepSeek V4 Flash GRPO support - Automodel path Sep 7, 2026
@zpqiu

zpqiu commented Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c7c4d76

@zpqiu
zpqiu marked this pull request as ready for review September 7, 2026 09:53
@zpqiu
zpqiu requested review from a team as code owners September 7, 2026 09:53
zpqiu and others added 18 commits September 10, 2026 04:59
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Preserve original refit errors, remove ineffective recipe settings, and
cover FP8 layout, tokenizer, and KV-cache contracts without expanding
model support.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
@zpqiu
zpqiu force-pushed the dsv4-automodel-r060-port branch from e874e79 to 0aa1b61 Compare September 10, 2026 13:04
Reject refit paths that bypass DeepSeek V4 preparation and finalization.
Require the nightly run to reach five steps with finite metrics, and pin
the validated attention backend.

Cover compatible FP32 scale references across two refits using real vLLM
quant configs. Document colocated H100 validation, precision mismatch,
Blackwell limitations, and the development branch's QAT work.

Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
@zpqiu

zpqiu commented Sep 10, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test d74d503

zpqiu and others added 2 commits September 11, 2026 06:37
Restore copy_ semantics for other FP8 MoE models and cover model selection,
parameter identity, and incompatible layouts in the refit tests.

Remove the Ray V1 executor override from the DeepSeek V4 recipe and guide.
Preserve upstream comments to keep the review diff focused.

Validation: 134 related unit tests and applicable pre-commit hooks passed
on a compute node. The RayExecutorV2 nightly completed five training steps
and passed all metric checks; its trailing pre-commit wrapper failed on
an unset SLURM_JOB_ID. Step 3 had a token_mult_prob_error spike of 1.795,
while the median remained within the configured threshold.

Signed-off-by: Zhaopeng Qiu <alexq@nvidia.com>
@sharonyu-115

Copy link
Copy Markdown
Contributor

/ok to test 628e0cc

@sharonyu-115 sharonyu-115 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks!

@sharonyu-115
sharonyu-115 enabled auto-merge (squash) September 14, 2026 07:43
@sharonyu-115
sharonyu-115 merged commit 8111533 into NVIDIA-NeMo:main Sep 14, 2026
147 of 150 checks passed
@yfw yfw mentioned this pull request Sep 18, 2026
4 tasks done
yfw added a commit that referenced this pull request Sep 18, 2026
…rwise reload

Two vLLM 0.29 changes silently broke the DeepSeek V4 refit path that #3595
validated on 0.25.1; both surfaced once 10e1a86 let the refit stream past
its first batch.

1. vLLM 0.29 gives DeepseekV4ForCausalLM a `packed_modules_mapping` in which
   `wkv` is a shard of two fused modules (`fused_wqa_wkv` and
   `fused_wkv_wgate`). `_get_module_from_param_name` resolved a shard by its
   leaf name alone, so `attn.wkv` was sent to `attn.fused_wkv_wgate`, the
   lookup failed (`Could not find module for parameter
   'model.layers.0.attn.wkv.weight'`), the weight was treated as non-fp8 and
   refit unquantized. Shard names that appear under more than one fused
   module are now left to the parent-aware DeepSeek V4 remap.

2. vLLM 0.29 split the layerwise-reload skip list: `SKIP_TENSORS` only keeps a
   tensor off the meta device, while the new `SKIP_LOAD_TENSORS` is what keeps
   its weight loader unwrapped and out of the layer's load accounting (0.25.1
   did both with `SKIP_TENSORS`). The immediately loaded expert tensors were
   therefore counted, vLLM ran `process_weights_after_loading` on each expert
   layer as soon as its last expert arrived, and `finalize_refit` converted
   the already kernel-layout weights and scales a second time. `prepare_refit`
   now registers the names in both sets and returns a `SkipNames` record so
   `restore_refit` removes exactly what was added from each set.

Also on the 0.29 test side: `WeightsMapper.get_unstacked_mapper` is
`get_rename_mapper` (what vLLM's own loader passes to `apply_vllm_mapper`),
`Fp8MoEMethod.__init__` reads the MoE config's partition sizes, the
import-target test skips imports inside `try`/`except ImportError` version
fallbacks, and the tied-alias filter tolerates module-less stand-ins (a bare
`load_weights` callable has no tied embeddings to drop).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants